# ====================================================================================
# 文件: preprocessing/step3_CoLaKG_item_graph.py
# 描述: [V4 格式修复版 - FAISS - 再次提供]
#      (核心) 修改了 torch.save 的逻辑，现在将
#      get_sparse_adj 返回的 edge_index 和 edge_weight
#      作为一个元组 (edge_index, edge_weight) 保存。
# ====================================================================================

import os
import torch
import numpy as np
import scipy.sparse as sp
from sklearn.preprocessing import normalize
from tqdm import tqdm
import faiss # 使用 FAISS 加速

# 导入共享的解析器和工具函数
from base_pre_parser import get_pre_config
# 导入 V5 版本的 get_sparse_adj (它返回 index, weight)
from utils import get_sparse_adj # 确保 utils.py 是 V5 版本

def main():
    print("--- [Step 3] CoLaKG: 物品语义图构建与归一化 (FAISS-GPU版) ---")
    config = get_pre_config()

    conf_data = config.data_config
    conf_pre = config.preproc_config.item_enhance
    DATA_DIR = config.DATA_DIR

    # 确定目标设备 (V34 逻辑)
    target_device = torch.device(
        "cuda:0" if torch.cuda.is_available() else "cpu"
    )
    faiss_gpu_index = 0 if torch.cuda.is_available() else -1

    print(f"[Step 3] 目标设备 (PyTorch): {target_device}")
    if faiss_gpu_index != -1:
        print(f"[Step 3] FAISS 将使用逻辑 GPU 索引: {faiss_gpu_index}")
    else:
        print("[Step 3] FAISS 将在 CPU 上运行 (未检测到 CUDA)。")

    # 1. 加载 SBERT 嵌入
    embed_file = os.path.join(DATA_DIR, conf_pre.output_embed_file)
    print(f"[Step 3.1] 加载 SBERT 嵌入: {embed_file}")

    if not os.path.exists(embed_file):
        print(f"错误: 物品嵌入文件未找到: {embed_file}")
        print("请先运行 step2_CoLaKG_item_embed.py。")
        return

    embeddings = torch.load(embed_file, map_location='cpu').numpy().astype(np.float32)
    n_items, d = embeddings.shape

    if n_items != conf_data.n_items:
        print(
            f"警告: 嵌入数量 ({n_items}) 与 config 中的 "
            f"n_items ({conf_data.n_items}) 不匹配。"
        )

    k = conf_pre.semantic_graph_k
    print(f"[Step 3.2] [FAISS 优化版] 开始计算 物品-物品 语义图 (K={k})")

    # [FAISS 优化]
    print("  > [FAISS] L2 归一化嵌入 (用于内积)...")
    embeddings_norm = normalize(embeddings, norm='l2', axis=1)

    print(f"  > [FAISS] 构建索引 (目标设备索引: {faiss_gpu_index})...")
    try:
        index_cpu = faiss.IndexFlatIP(d)
        if faiss_gpu_index != -1:
            res = faiss.StandardGpuResources()
            gpu_index = faiss.index_cpu_to_gpu(res, faiss_gpu_index, index_cpu)
            gpu_index.add(embeddings_norm)
            index_to_search = gpu_index
            print("  > [FAISS] 使用 GPU 索引进行搜索。")
        else:
            index_cpu.add(embeddings_norm)
            index_to_search = index_cpu
            print("  > [FAISS] 使用 CPU 索引进行搜索。")
    except Exception as e:
        print(f"❌ 错误: FAISS 索引失败: {e}")
        if "faiss.StandardGpuResources" in str(e):
             print("请确保 'faiss-gpu' 已正确安装 (pip install faiss-gpu)。")
        return

    print(f"  > [FAISS] 搜索 Top-{k+1} 邻居...")
    D, I = index_to_search.search(embeddings_norm, k + 1)
    print("  > [FAISS] 搜索完成。")

    # 3. 构建 Scipy 稀疏邻接矩阵 (在 CPU 上)
    print("[Step 3.3] 构建稀疏邻接矩阵 (Scipy)...")
    row, col, data = [], [], []
    for i in tqdm(range(n_items), desc="  > 构建图中"):
        for j_idx in range(k + 1):
            neighbor_id = I[i, j_idx]
            if neighbor_id == i or neighbor_id == -1:
                continue
            if 0 <= neighbor_id < n_items:
                row.append(i)
                col.append(neighbor_id)
                data.append(D[i, j_idx])
            # else: # 可选：打印无效邻居警告
            #      print(f"[Step 3] 警告: FAISS 返回了无效邻居索引 {neighbor_id}")

    adj = sp.csr_matrix((data, (row, col)), shape=(n_items, n_items))
    adj = adj.maximum(adj.T) # 确保对称
    print(f"  > 稀疏图构建完毕。边数 (对称): {adj.nnz}")

    # 4. 图归一化 (使用 utils.py V5)
    print(f"[Step 3.4] 开始图归一化 (目标设备: {target_device})...")
    # !! get_sparse_adj V5 返回 edge_index, edge_weight !!
    norm_ii_edge_index, norm_ii_edge_weight = get_sparse_adj(
        adj, device=target_device, gmae_norm=False # 使用 LightGCN 归一化
    )

    # 5. 保存 (核心修复：保存元组)
    output_file = os.path.join(DATA_DIR, conf_pre.output_graph_file)
    print(f"[Step 3.5] 保存归一化图 (edge_index, edge_weight) 到: {output_file}")
    # !! 保存为一个包含两个张量的元组 (确保张量在 CPU 上保存) !!
    torch.save(
        (norm_ii_edge_index.cpu(), norm_ii_edge_weight.cpu()),
        output_file
    )

    print(f"\n✅ [Step 3] 成功完成!")
    print(f"归一化物品语义图 (edge_index, edge_weight) 已保存到: {output_file}")


if __name__ == "__main__":
    main()